import { app } from "../../../scripts/app.js";

class PoseEditorWidget {
    constructor(node) {
        this.node = node;
        this.canvas = null;
        this.ctx = null;
        this.poses = [];
        this.backgroundImage = null;
        this.isEditing = false;
        this.selectedPoint = null;
        this.dragOffset = { x: 0, y: 0 };
        this.history = [];
        this.historyIndex = -1;
        this.gridSnap = false;
        this.symmetryMode = false;
        
        // COCO 17个关键点的连接关系
        this.connections = [
            [0, 1], [0, 2], [1, 3], [2, 4],      // 头部
            [5, 6],                               // 肩膀
            [5, 7], [7, 9],                       // 左臂
            [6, 8], [8, 10],                      // 右臂
            [5, 11], [6, 12],                     // 肩膀到臀部
            [11, 12],                             // 臀部
            [11, 13], [13, 15],                   // 左腿
            [12, 14], [14, 16]                    // 右腿
        ];
        
        // 关键点名称
        this.keypointNames = [
            'nose', 'left_eye', 'right_eye', 'left_ear', 'right_ear',
            'left_shoulder', 'right_shoulder', 'left_elbow', 'right_elbow',
            'left_wrist', 'right_wrist', 'left_hip', 'right_hip',
            'left_knee', 'right_knee', 'left_ankle', 'right_ankle'
        ];
        
        this.setupCanvas();
        this.bindEvents();
    }
    
    setupCanvas() {
        // 创建编辑画布容器
        const container = document.createElement('div');
        container.style.cssText = `
            position: relative;
            width: 100%;
            max-width: 600px;
            margin: 10px 0;
            border: 2px solid #333;
            border-radius: 8px;
            background: #1e1e1e;
            overflow: hidden;
        `;
        
        // 创建头部工具栏
        const toolbar = document.createElement('div');
        toolbar.style.cssText = `
            background: #2a2a2a;
            padding: 8px;
            display: flex;
            gap: 10px;
            align-items: center;
            font-size: 12px;
            color: #ccc;
        `;
        
        // 编辑模式切换
        const editBtn = document.createElement('button');
        editBtn.textContent = '✏️ 开始编辑';
        editBtn.style.cssText = `
            background: #4a5568;
            color: white;
            border: none;
            padding: 4px 8px;
            border-radius: 4px;
            cursor: pointer;
            font-size: 11px;
        `;
        editBtn.onclick = () => this.toggleEditMode();
        
        // 重置按钮
        const resetBtn = document.createElement('button');
        resetBtn.textContent = '🔄 重置';
        resetBtn.style.cssText = editBtn.style.cssText;
        resetBtn.onclick = () => this.resetPose();
        
        // 撤销按钮
        const undoBtn = document.createElement('button');
        undoBtn.textContent = '↶ 撤销';
        undoBtn.style.cssText = editBtn.style.cssText;
        undoBtn.onclick = () => this.undo();
        
        // 重做按钮
        const redoBtn = document.createElement('button');
        redoBtn.textContent = '↷ 重做';
        redoBtn.style.cssText = editBtn.style.cssText;
        redoBtn.onclick = () => this.redo();
        
        // 网格对齐开关
        const gridToggle = document.createElement('input');
        gridToggle.type = 'checkbox';
        gridToggle.id = 'gridSnap';
        gridToggle.onchange = (e) => {
            this.gridSnap = e.target.checked;
            this.redraw();
        };
        
        const gridLabel = document.createElement('label');
        gridLabel.htmlFor = 'gridSnap';
        gridLabel.textContent = '📐 网格对齐';
        gridLabel.style.fontSize = '11px';
        
        // 对称编辑开关
        const symmetryToggle = document.createElement('input');
        symmetryToggle.type = 'checkbox';
        symmetryToggle.id = 'symmetryMode';
        symmetryToggle.onchange = (e) => {
            this.symmetryMode = e.target.checked;
        };
        
        const symmetryLabel = document.createElement('label');
        symmetryLabel.htmlFor = 'symmetryMode';
        symmetryLabel.textContent = '🔄 对称编辑';
        symmetryLabel.style.fontSize = '11px';
        
        toolbar.appendChild(editBtn);
        toolbar.appendChild(resetBtn);
        toolbar.appendChild(undoBtn);
        toolbar.appendChild(redoBtn);
        toolbar.appendChild(document.createTextNode(' | '));
        toolbar.appendChild(gridToggle);
        toolbar.appendChild(gridLabel);
        toolbar.appendChild(document.createTextNode(' '));
        toolbar.appendChild(symmetryToggle);
        toolbar.appendChild(symmetryLabel);
        
        // 创建画布
        this.canvas = document.createElement('canvas');
        this.canvas.width = 512;
        this.canvas.height = 512;
        this.canvas.style.cssText = `
            display: block;
            background: #000;
            cursor: crosshair;
        `;
        
        this.ctx = this.canvas.getContext('2d');
        
        // 状态显示
        const statusDiv = document.createElement('div');
        statusDiv.style.cssText = `
            background: #2a2a2a;
            padding: 6px 8px;
            font-size: 11px;
            color: #999;
        `;
        statusDiv.textContent = '💡 点击"开始编辑"后可拖拽蓝色圆点调整关键点位置';
        this.statusDiv = statusDiv;
        
        container.appendChild(toolbar);
        container.appendChild(this.canvas);
        container.appendChild(statusDiv);
        
        // 添加到节点
        const widget = this.node.addDOMWidget("pose_editor", "div", container);
        this.widget = widget;
        
        this.redraw();
    }
    
    bindEvents() {
        // 鼠标事件
        this.canvas.addEventListener('mousedown', this.onMouseDown.bind(this));
        this.canvas.addEventListener('mousemove', this.onMouseMove.bind(this));
        this.canvas.addEventListener('mouseup', this.onMouseUp.bind(this));
        this.canvas.addEventListener('mouseleave', this.onMouseUp.bind(this));
        
        // 键盘事件（需要聚焦）
        this.canvas.setAttribute('tabindex', '0');
        this.canvas.addEventListener('keydown', this.onKeyDown.bind(this));
        
        // 禁用右键菜单
        this.canvas.addEventListener('contextmenu', (e) => e.preventDefault());
    }
    
    onMouseDown(e) {
        if (!this.isEditing) return;
        
        const rect = this.canvas.getBoundingClientRect();
        const x = e.clientX - rect.left;
        const y = e.clientY - rect.top;
        
        // 查找最近的关键点
        this.selectedPoint = this.findNearestKeypoint(x, y);
        
        if (this.selectedPoint) {
            const kp = this.selectedPoint.keypoint;
            this.dragOffset.x = x - kp.x;
            this.dragOffset.y = y - kp.y;
            
            // 保存历史状态
            this.saveHistory();
            
            this.canvas.style.cursor = 'grabbing';
            this.updateStatus(`🎯 编辑关键点: ${this.keypointNames[this.selectedPoint.index]} (#${this.selectedPoint.index})`);
        }
    }
    
    onMouseMove(e) {
        if (!this.isEditing) return;
        
        const rect = this.canvas.getBoundingClientRect();
        const x = e.clientX - rect.left;
        const y = e.clientY - rect.top;
        
        if (this.selectedPoint) {
            // 拖拽移动关键点
            let newX = x - this.dragOffset.x;
            let newY = y - this.dragOffset.y;
            
            // 网格对齐
            if (this.gridSnap) {
                const gridSize = 20;
                newX = Math.round(newX / gridSize) * gridSize;
                newY = Math.round(newY / gridSize) * gridSize;
            }
            
            // 边界检查
            newX = Math.max(0, Math.min(this.canvas.width - 1, newX));
            newY = Math.max(0, Math.min(this.canvas.height - 1, newY));
            
            this.selectedPoint.keypoint.x = newX;
            this.selectedPoint.keypoint.y = newY;
            
            // 对称编辑模式
            if (this.symmetryMode) {
                this.applySymmetricEdit(this.selectedPoint.index, newX, newY);
            }
            
            this.redraw();
            this.updateNodeData();
            
        } else {
            // 鼠标悬停检测
            const nearPoint = this.findNearestKeypoint(x, y);
            if (nearPoint) {
                this.canvas.style.cursor = 'grab';
                this.updateStatus(`📍 ${this.keypointNames[nearPoint.index]} (#${nearPoint.index}) - 点击拖拽编辑`);
            } else {
                this.canvas.style.cursor = 'crosshair';
                this.updateStatus('🖱️ 拖拽蓝色圆点编辑关键点位置');
            }
        }
    }
    
    onMouseUp(e) {
        if (this.selectedPoint) {
            this.selectedPoint = null;
            this.canvas.style.cursor = 'crosshair';
            this.updateStatus('✅ 关键点位置已更新');
        }
    }
    
    onKeyDown(e) {
        if (!this.isEditing) return;
        
        if (e.ctrlKey || e.metaKey) {
            if (e.key === 'z' && !e.shiftKey) {
                e.preventDefault();
                this.undo();
            } else if (e.key === 'y' || (e.key === 'z' && e.shiftKey)) {
                e.preventDefault();
                this.redo();
            }
        }
    }
    
    findNearestKeypoint(x, y) {
        const threshold = 15; // 点击检测范围
        let nearest = null;
        let minDistance = threshold;
        
        for (let poseIndex = 0; poseIndex < this.poses.length; poseIndex++) {
            const pose = this.poses[poseIndex];
            
            for (let i = 0; i < pose.keypoints.length; i++) {
                const kp = pose.keypoints[i];
                
                if (kp.confidence > 0) {
                    const distance = Math.sqrt((x - kp.x) ** 2 + (y - kp.y) ** 2);
                    
                    if (distance < minDistance) {
                        minDistance = distance;
                        nearest = {
                            poseIndex: poseIndex,
                            index: i,
                            keypoint: kp
                        };
                    }
                }
            }
        }
        
        return nearest;
    }
    
    applySymmetricEdit(keypointIndex, newX, newY) {
        // 实现左右对称编辑
        const symmetricPairs = [
            [1, 2], [3, 4],                   // 眼睛、耳朵
            [5, 6], [7, 8], [9, 10],         // 肩膀、肘部、手腕
            [11, 12], [13, 14], [15, 16]     // 臀部、膝盖、脚踝
        ];
        
        for (const [left, right] of symmetricPairs) {
            if (keypointIndex === left) {
                // 编辑左侧，同步到右侧
                const symmetricX = this.canvas.width - newX;
                if (this.poses[0] && this.poses[0].keypoints[right]) {
                    this.poses[0].keypoints[right].x = symmetricX;
                    this.poses[0].keypoints[right].y = newY;
                }
                break;
            } else if (keypointIndex === right) {
                // 编辑右侧，同步到左侧
                const symmetricX = this.canvas.width - newX;
                if (this.poses[0] && this.poses[0].keypoints[left]) {
                    this.poses[0].keypoints[left].x = symmetricX;
                    this.poses[0].keypoints[left].y = newY;
                }
                break;
            }
        }
    }
    
    toggleEditMode() {
        this.isEditing = !this.isEditing;
        const btn = this.widget.element.querySelector('button');
        
        if (this.isEditing) {
            btn.textContent = '💾 完成编辑';
            btn.style.background = '#38a169';
            this.updateStatus('🎯 编辑模式已启用 - 拖拽蓝色圆点调整位置');
        } else {
            btn.textContent = '✏️ 开始编辑';
            btn.style.background = '#4a5568';
            this.updateStatus('💡 点击"开始编辑"后可拖拽蓝色圆点调整关键点位置');
        }
        
        this.canvas.style.cursor = this.isEditing ? 'crosshair' : 'default';
        this.redraw();
    }
    
    saveHistory() {
        // 保存当前状态到历史记录
        const state = JSON.parse(JSON.stringify(this.poses));
        this.history = this.history.slice(0, this.historyIndex + 1);
        this.history.push(state);
        this.historyIndex++;
        
        // 限制历史记录数量
        if (this.history.length > 20) {
            this.history.shift();
            this.historyIndex--;
        }
    }
    
    undo() {
        if (this.historyIndex > 0) {
            this.historyIndex--;
            this.poses = JSON.parse(JSON.stringify(this.history[this.historyIndex]));
            this.redraw();
            this.updateNodeData();
            this.updateStatus('↶ 已撤销上一步操作');
        }
    }
    
    redo() {
        if (this.historyIndex < this.history.length - 1) {
            this.historyIndex++;
            this.poses = JSON.parse(JSON.stringify(this.history[this.historyIndex]));
            this.redraw();
            this.updateNodeData();
            this.updateStatus('↷ 已重做操作');
        }
    }
    
    resetPose() {
        if (confirm('确定要重置骨架到默认姿态吗？')) {
            this.poses = this.createDefaultPose();
            this.saveHistory();
            this.redraw();
            this.updateNodeData();
            this.updateStatus('🔄 已重置为默认姿态');
        }
    }
    
    createDefaultPose() {
        // 创建默认T字形姿态
        const centerX = this.canvas.width / 2;
        const centerY = this.canvas.height / 2;
        const scale = Math.min(this.canvas.width, this.canvas.height) * 0.15;
        
        const defaultOffsets = [
            [0, -0.15],      // nose
            [-0.02, -0.12],  // left_eye
            [0.02, -0.12],   // right_eye
            [-0.04, -0.10],  // left_ear
            [0.04, -0.10],   // right_ear
            [-0.20, -0.05],  // left_shoulder
            [0.20, -0.05],   // right_shoulder
            [-0.35, 0.05],   // left_elbow
            [0.35, 0.05],    // right_elbow
            [-0.50, 0.15],   // left_wrist
            [0.50, 0.15],    // right_wrist
            [-0.10, 0.25],   // left_hip
            [0.10, 0.25],    // right_hip
            [-0.10, 0.50],   // left_knee
            [0.10, 0.50],    // right_knee
            [-0.10, 0.75],   // left_ankle
            [0.10, 0.75],    // right_ankle
        ];
        
        const keypoints = defaultOffsets.map(([dx, dy]) => ({
            x: centerX + dx * scale,
            y: centerY + dy * scale,
            confidence: 1.0
        }));
        
        return [{
            id: 0,
            keypoints: keypoints,
            overall_confidence: 1.0
        }];
    }
    
    redraw() {
        // 清空画布
        this.ctx.fillStyle = '#000';
        this.ctx.fillRect(0, 0, this.canvas.width, this.canvas.height);
        
        // 绘制背景图像
        if (this.backgroundImage) {
            this.ctx.globalAlpha = 0.7;
            this.ctx.drawImage(this.backgroundImage, 0, 0, this.canvas.width, this.canvas.height);
            this.ctx.globalAlpha = 1.0;
        }
        
        // 绘制网格
        if (this.gridSnap) {
            this.drawGrid();
        }
        
        // 绘制骨架
        this.drawSkeleton();
        
        // 绘制编辑手柄
        if (this.isEditing) {
            this.drawEditHandles();
        }
    }
    
    drawGrid() {
        const gridSize = 20;
        this.ctx.strokeStyle = '#333';
        this.ctx.lineWidth = 1;
        this.ctx.setLineDash([2, 2]);
        
        // 垂直线
        for (let x = 0; x <= this.canvas.width; x += gridSize) {
            this.ctx.beginPath();
            this.ctx.moveTo(x, 0);
            this.ctx.lineTo(x, this.canvas.height);
            this.ctx.stroke();
        }
        
        // 水平线
        for (let y = 0; y <= this.canvas.height; y += gridSize) {
            this.ctx.beginPath();
            this.ctx.moveTo(0, y);
            this.ctx.lineTo(this.canvas.width, y);
            this.ctx.stroke();
        }
        
        this.ctx.setLineDash([]);
    }
    
    drawSkeleton() {
        const colors = ['#00ff00', '#ff0066', '#0066ff', '#ffff00', '#ff6600'];
        
        for (let poseIndex = 0; poseIndex < this.poses.length; poseIndex++) {
            const pose = this.poses[poseIndex];
            const color = colors[poseIndex % colors.length];
            
            // 绘制连接线
            this.ctx.strokeStyle = color;
            this.ctx.lineWidth = 2;
            
            for (const [start, end] of this.connections) {
                const startKp = pose.keypoints[start];
                const endKp = pose.keypoints[end];
                
                if (startKp.confidence > 0 && endKp.confidence > 0) {
                    this.ctx.beginPath();
                    this.ctx.moveTo(startKp.x, startKp.y);
                    this.ctx.lineTo(endKp.x, endKp.y);
                    this.ctx.stroke();
                }
            }
        }
    }
    
    drawEditHandles() {
        for (let poseIndex = 0; poseIndex < this.poses.length; poseIndex++) {
            const pose = this.poses[poseIndex];
            
            for (let i = 0; i < pose.keypoints.length; i++) {
                const kp = pose.keypoints[i];
                
                if (kp.confidence > 0) {
                    // 绘制编辑手柄
                    this.ctx.fillStyle = '#fff';
                    this.ctx.strokeStyle = '#007acc';
                    this.ctx.lineWidth = 2;
                    
                    this.ctx.beginPath();
                    this.ctx.arc(kp.x, kp.y, 8, 0, 2 * Math.PI);
                    this.ctx.fill();
                    this.ctx.stroke();
                    
                    // 绘制编号
                    this.ctx.fillStyle = '#000';
                    this.ctx.font = '10px Arial';
                    this.ctx.textAlign = 'center';
                    this.ctx.fillText(i.toString(), kp.x, kp.y + 3);
                }
            }
        }
    }
    
    updateStatus(message) {
        if (this.statusDiv) {
            this.statusDiv.textContent = message;
        }
    }
    
    updateNodeData() {
        // 更新节点的姿态数据
        if (this.poses.length > 0) {
            const jsonData = JSON.stringify({
                timestamp: new Date().toISOString(),
                image_size: {
                    width: this.canvas.width,
                    height: this.canvas.height
                },
                persons: this.poses.map(pose => ({
                    id: pose.id,
                    keypoints: pose.keypoints.map((kp, i) => ({
                        name: this.keypointNames[i],
                        x: kp.x,
                        y: kp.y,
                        confidence: kp.confidence
                    })),
                    overall_confidence: pose.overall_confidence
                }))
            }, null, 2);
            
            // 查找姿态数据输入框并更新
            const poseDataWidget = this.node.widgets.find(w => w.name === 'pose_data');
            if (poseDataWidget) {
                poseDataWidget.value = jsonData;
            }
        }
    }
    
    loadPoseData(jsonString) {
        try {
            if (!jsonString || jsonString.trim() === '') {
                this.poses = this.createDefaultPose();
                this.redraw();
                return;
            }
            
            const data = JSON.parse(jsonString);
            
            if (data.persons && Array.isArray(data.persons)) {
                this.poses = data.persons.map((person, index) => ({
                    id: index,
                    keypoints: person.keypoints || [],
                    overall_confidence: person.overall_confidence || 1.0
                }));
                
                // 确保每个姿态有17个关键点
                this.poses.forEach(pose => {
                    while (pose.keypoints.length < 17) {
                        pose.keypoints.push({ x: 0, y: 0, confidence: 0 });
                    }
                    pose.keypoints = pose.keypoints.slice(0, 17);
                });
                
                this.saveHistory();
                this.redraw();
                this.updateStatus(`📥 已导入 ${this.poses.length} 个姿态`);
            }
        } catch (e) {
            console.error('解析姿态数据失败:', e);
            this.updateStatus('❌ 姿态数据格式错误');
        }
    }
    
    setBackgroundImage(imageUrl) {
        if (imageUrl) {
            const img = new Image();
            img.onload = () => {
                this.backgroundImage = img;
                this.redraw();
            };
            img.src = imageUrl;
        }
    }
}

// 注册节点扩展
app.registerExtension({
    name: "ComfyUI-xhh.PoseEditor",
    
    async beforeRegisterNodeDef(nodeType, nodeData, app) {
        if (nodeData.name === "XHH_PoseEditor") {
            
            const onNodeCreated = nodeType.prototype.onNodeCreated;
            nodeType.prototype.onNodeCreated = function () {
                const result = onNodeCreated?.apply(this, arguments);
                
                // 添加编辑器实例
                this.poseEditor = new PoseEditorWidget(this);
                
                // 监听姿态数据输入变化
                const poseDataWidget = this.widgets.find(w => w.name === 'pose_data');
                if (poseDataWidget) {
                    const originalCallback = poseDataWidget.callback;
                    poseDataWidget.callback = (value) => {
                        if (originalCallback) originalCallback(value);
                        this.poseEditor.loadPoseData(value);
                    };
                }
                
                return result;
            };
            
            // 自定义节点大小
            const onDrawBackground = nodeType.prototype.onDrawBackground;
            nodeType.prototype.onDrawBackground = function(ctx) {
                if (onDrawBackground) {
                    onDrawBackground.apply(this, arguments);
                }
                
                // 确保节点足够大以容纳编辑器
                if (this.size[1] < 600) {
                    this.size[1] = 650;
                }
            };
        }
    }
}); 